DB = 'db4coco'

if '__main__' == __name__:
    import os
    from PyCmpltrtok.common import sep
    from PyCmpltrtok.common_pymongo import PymongoDict
    import json

    CAPTION_OK = True
    IMAGE_OK = True
    CATEGORY_OK = False

    def _main():
        # load captions
        if not CAPTION_OK:
            json_path_cap = '/home/asuspei/large_data/_very_large/coco2014/annotations/captions_val2014.json'
            sep(json_path_cap)
            with open(json_path_cap, 'r', encoding='utf8') as f:
                json_obj_cap = json.load(f)
            sep('CAPTION')
            id2cap = PymongoDict(DB, 'id2cap')
            cnt = 0
            for xdict in json_obj_cap['annotations']:
                cnt += 1
                if cnt % 100 == 0:
                    print(cnt)
                xid = xdict['image_id']
                xcaption = xdict['caption']
                id2cap[xid] = xcaption

        # load instants
        json_path_inst = '/home/asuspei/large_data/_very_large/coco2014/annotations/instances_val2014.json'
        sep(json_path_inst)
        with open(json_path_inst, 'r', encoding='utf8') as f:
            json_obj_inst = json.load(f)

        if not IMAGE_OK:
            sep('IMAGES')
            ids = []
            id2name, id2wh = PymongoDict(DB, 'id2name'), PymongoDict(DB, 'id2wh')
            cnt = 0
            for xdict in json_obj_inst['images']:
                cnt += 1
                if cnt % 100 == 0:
                    print(cnt)
                xid = xdict['id']
                ids.append(xid)
                xname = xdict['file_name']
                id2name[xid] = xname
                xw = xdict['width']
                xh = xdict['height']
                id2wh[xid] = [xw, xh]
            ids = sorted(ids)
            lists = PymongoDict(DB, 'lists')
            lists['ids'] = ids

        if not CATEGORY_OK:
            sep('CATEGORIES')
            cate_id2cate, cate_id2cate_super = PymongoDict(DB, 'cate_id2cate'), PymongoDict(DB, 'cate_id2cate_super')
            cnt = 0
            for xdict in json_obj_inst['categories']:
                cnt += 1
                if cnt % 100 == 0:
                    print(cnt)
                cate_id = xdict['id']
                xcate = xdict['name']
                xcate_super = xdict['supercategory']
                cate_id2cate[cate_id] = xcate
                cate_id2cate_super[cate_id] = xcate_super

        sep('ANNOTATIONS')
        id2ann_ids = PymongoDict(DB, 'id2ann_ids')  # image id to annotations ids
        ann_id2cate_id, ann_id2bbox, ann_id2seg = PymongoDict(DB, 'ann_id2cate_id'), PymongoDict(DB, 'ann_id2bbox'), PymongoDict(DB, 'ann_id2seg')
        cnt = 0
        for xdict in json_obj_inst['annotations']:
            cnt += 1
            if cnt % 100 == 0:
                print(cnt)
            ann_id = xdict['id']
            xid = xdict['image_id']
            xlist = id2ann_ids[xid]
            if xlist is None:
                xlist = []
            xlist.append(ann_id)
            id2ann_ids[xid] = xlist
            cate_id = xdict['category_id']
            ann_id2cate_id[ann_id] = cate_id
            xbbox = xdict['bbox']
            ann_id2bbox[ann_id] = xbbox
            xseg = xdict['segmentation']
            ann_id2seg[ann_id] = xseg

    _main()
